from global_val import GlobalVal
import covar
import math
from cm import *

gv = GlobalVal('./gv_data/GlobalVal.pickle')
MAX_PRIO = 10


def decide_what_to_do_next(state, flag):
    # 记得判断自己！
    # 返回值种类：
    # ("info","使用手法的名称",可选参数列表)
    # ("attack","攻击手法的名称",可选参数列表)
    # ("move",移动的地址)
    # ("nothing")

    # 在转移后立即更新topo
    # 注意：什么时候等待？——自己的order序号不是最小值（判断自己的order序号取值情况后）+其他主机更新order时（控制权转移，对应move转变）

    # next_step_ip = eval(gv.get_val("order", []))[covar.get_val("nextStep")][0]


    # if not isNeighborAndFirst(next_step_ip, eval(gv.get_val("topo", {}))):
    #     return ["wait"] ####
    # else:
    #     return ["attack"]

    # if state == "wait":
    #     nextStep = mergeSort(eval(gv.get_val("order", [])))
    #     covar.set_val("nextStep", nextStep)
    #     next_step_ip = eval(gv.get_val("order", []))[covar.get_val("nextStep")][0]
    #     next_step_leak = eval(gv.get_val("order", []))[covar.get_val("nextStep")][1]
    #
    # else:
    #     1
    #
    #
    # nextStep = mergeSort(eval(gv.get_val("order", [])))
    # covar.set_val("nextStep", nextStep)
    # if state == "info4use":  # 传输已攻击下来的主机的有用信息——info4attack
    #     return ["info4attack"]
    # elif state == "info4attack":  # 扫描端口信息，更新全局变量：order
    #     # 如果端口不匹配：order序号+1——扫描下一个更新序号对应的漏洞端口info4attack
    #     if not is_matched_port:
    #         order = eval(gv.get_val("order"))
    #         order[covar.get_val("next_order_index")][1] += 1
    #         gv.set_val("order", str(order))
    #         next_bug_index, next_attack_ip = mergeSort(eval(gv.get_val("order", [])))
    #         covar.set_val("next_bug_index", next_bug_index)
    #         covar.set_val("next_attack_ip", next_attack_ip)
    #         return ["info4attack", next_bug_index]
    #     # 如果端口匹配：如果（自己与下一个攻击ip临近且是attacked_ip_list中的第一个ip）或者是第一个节点，则attack；否则wait
    #     else:
    #         if isNeighborAndFirst(next_step_ip, eval(gv.get_val("topo", {}))):
    #             return ["attack", info]
    #         else:
    #             return ["info4attack", "wait"]
    # elif state == "attack":  # 根据order序号匹配att_met中的漏洞，得到漏洞名称和端口，调用攻击脚本进行攻击；更新全局变量：order
    #     # 如果验证攻击失败：order序号+1——扫描下一个更新序号对应的漏洞端口info4attack
    #     if failed_attack:
    #
    #         return
    #     else:
    #     # 如果验证攻击成功：order序号设为最大值，正式实施攻击进行转移——move
    # elif state == "move":  # 转移后进入info4attack状态等待order更新，激活主线程
    #     return ["wait"]

    #####################
    if state == "info4use" and flag is False:  # 传输已攻击下来的主机的有用信息——info4attack
        return ["info4use"]
    elif (state == "info4attack" and flag is False) or (state == "attack" and flag is False) \
            or (state == "info4use" and flag is True):
        # 下一个扫描的主机为序号-1时说明刚初始化,否则更新上次扫描失败后的order
        if state == "info4attack" and covar.get_val("screen_failed") is True:
            tmp_order = eval(gv.get_val("order"))
            tmp_order[covar.get_val("nextStep")][1] += 1
            print("22222", tmp_order)
            gv.set_val("order", str(tmp_order))
            covar.set_val("screen_failed", False)

        # 排序获得最优步序号 getMin
        minOrder = mergeSort(eval(gv.get_val("order", [])))
        covar.set_val("minOrder", minOrder)

        # 顺序扫描获得第一个可能最优步
        for i in range(len(eval(gv.get_val("order", [])))):
            if eval(gv.get_val("order", []))[i][1] == minOrder:
                covar.set_val("nextStep", i)
                break

        # 后续使用：下一步扫描的主机ip和漏洞序号
        next_step_ip = eval(gv.get_val("order", []))[covar.get_val("nextStep")][0]
        if state == "info4attack" and covar.get_val("screen_failed") is True:
            next_step_leak = eval(gv.get_val("order", []))[covar.get_val("nextStep")][1] - 1
        else:
            next_step_leak = eval(gv.get_val("order", []))[covar.get_val("nextStep")][1]

        # 检查下一步是否为为当前主机的相邻节点 isNeighbor 如果是，本机扫描，如果不是，order需要更新则广播，不需要则等待
        # 决定计算是否由自己控制时，不仅要求与待攻击节点相邻，而且需要一次从attacked_ips_list中取出ip，如果自己是第一个满足要求的才能有控制权，否则没有
        print("next_step_ip", next_step_ip)
        if isNeighborAndFirst(next_step_ip, eval(gv.get_val("topo", {}))):
            covar.set_val("updatedOrder", True)
            print("1234566666666", next_step_leak,covar.get_val("att_met"))
            port = int(covar.get_val("att_met")[next_step_leak - 1][0]) # changed
            return ["info4attack", next_step_ip, port]
        else:
            if covar.get_val("updatedOrder") is True:
                # 广播只在本地更新后的order,重设order更新标志
                timestamp = str(time.time())
                gv.set_val("timestamp", timestamp)
                # broadcast_wakeup(timestamp)
                broadcast_neworder(eval(gv.get_val("order", [])), timestamp)
                covar.set_val("updatedOrder", False)
            return ["info4attack", "wait"]

    elif state == "info4attack" and flag is True:
        # 更新上次扫描成功后的order,因为有可能攻击失败再用到order______________????????不能攻击某漏洞（包括扫描端口不匹配/攻击失败）才应该order序号+1
        tmp_order = eval(gv.get_val("order", []))
        tmp_order[covar.get_val("nextStep")][1] += 1
        gv.set_val("order", str(tmp_order))

        # 后续使用：下一步扫描的主机ip和漏洞序号
        next_step_ip = eval(gv.get_val("order", []))[covar.get_val("nextStep")][0]
        next_step_leak = eval(gv.get_val("order", []))[covar.get_val("nextStep")][1] - 1
        port = int(covar.get_val("att_met")[next_step_leak][0])
        attack_method = covar.get_val("att_met")[next_step_leak][1]
        return ["attack", next_step_ip, port, attack_method]

    elif state == "attack" and flag is True:
        next_step_to_move_ip = eval(gv.get_val("order", []))[covar.get_val("nextStep")][0]
        next_step_leak = eval(gv.get_val("order", []))[covar.get_val("nextStep")][1] - 1
        next_step_to_move_port = int(covar.get_val("att_met")[next_step_leak][0])
        move_method = covar.get_val("att_met")[next_step_leak][1]
        return ["move", next_step_to_move_ip, next_step_to_move_port, move_method]

    elif state == "move":
        return ["info4attack", "wait"]

# 得到order中序号最小的行对应的index
def mergeSort(order_list):
    next_bug_index = MAX_PRIO
    next_attack_ip = ''
    for order in order_list:
        if order[1] < next_bug_index:
            next_bug_index = order[1]
            # next_attack_ip = order[0]
    # return next_bug_index, next_attack_ip
    return next_bug_index
    # for i in range(len(order_list)):
    #     if next_bug_ID == order_list[i][1]:
    #         return i

# 计算是否由自己控制时，不仅要求与待攻击节点相邻，而且需要一次从attacked_ips_list中取出ip，如果自己是第一个满足要求的才能有控制权，否则没有
def isNeighborAndFirst(next_attack_ip, topo_dict):
    attacked_ips_list = eval(gv.get_val("attacked_ips_list", []))
    control_ip = ''
    for ip in attacked_ips_list:
        print(ip)
        if next_attack_ip in topo_dict[ip]:
            control_ip = ip
            break
    self_ip = covar.get_val("self_ip")
    isFirst = gv.get_val("isFirst")
    if self_ip == control_ip or isFirst == 'True':
        return True
    else:
        return False

    # self_ip = covar.get_val("self_ip")
    # for ip in self_ip:
    #     if nextStepIP in topo_dict[ip]:
    #         return True
    # return False
